from unittest import IsolatedAsyncioTestCase
from benchmark import ClsBenchmark, BarBenchmark, LineBenchmark, PieBenchmark

from config import generate_dataset_label_path, model_predict_data, generate_dataset_images_path
import os

class pictogram_benchmark_test(IsolatedAsyncioTestCase):

    # 测试公共元素识别率
    def test_cls(self):
        # yolo result
        yolo_pred_path = os.path.join(model_predict_data, "YoloCls.json")
        # corner result
        corner_pred_path = os.path.join(model_predict_data, "CornerCls.json")
        # cached result
        cached_pred_path = os.path.join(model_predict_data, "CachedCls.json")
        cls_benchmark = ClsBenchmark()
        cls_benchmark.load(yolo_pred_path, corner_pred_path, cached_pred_path, generate_dataset_label_path)
        cls_benchmark.eval()

    # 测试条形图数据元素
    def test_bar(self):
        bar_pred_path = os.path.join(model_predict_data, "BarDetect.json")
        bar_benchmark = BarBenchmark()
        bar_benchmark.load(bar_pred_path, generate_dataset_label_path)
        bar_benchmark.eval(IOU_threshold=0.78)
        pass

    # 测试折线图数据元素
    def test_line(self):
        line_pred_path = os.path.join(model_predict_data, "Lineformer.json")
        benchmark = LineBenchmark(generate_dataset_images_path, generate_dataset_label_path, line_pred_path)
        benchmark.run()
        benchmark.summary()

    # 测试饼图数据元素
    def test_pie(self):
        pie_pred_path = os.path.join(model_predict_data, "PieDetect.json")
        benchmark = PieBenchmark(generate_dataset_label_path, pie_pred_path)
        benchmark.run(threshold=0.15)
        benchmark.summary()
